# Dickstein, Ho, and Mark (2023)
# This script creates data used in table 5 of the paper

# Preliminaries
setwd("../../../../library")
source("PreliminariesCode.R")
library(readr)
library(knitr)
library(readxl)
library(numDeriv)
library(stringr)
library(yaml)
library(haven)
try(library(Hmisc))

main <- function() {
  get_derived("lowrisk_simplemh_censor0.2", "main")
  get_derived("lowrisk_simplemh_censor0.5", "main")
  
  get_derived("lowrisk_simplemh_censor0.2", "switcher")
  get_derived("lowrisk_simplemh_censor0.5", "switcher")
}

get_derived <- function(run, type) {
  
  output_path <- paste0(project_folder, "/analysis/tablesandfigures/release/derived_paramests/specs/", run, "_", type)
  # Prepare file path
  dir.create(output_path, recursive = T)
  f <- list.files(output_path, include.dirs = F, full.names = T)
  file.remove(f)
  
  sink(paste0(output_path, "/log.txt"))
  
  # Load the actual estimates
    covar_names_csv <- read_csv(paste0(project_folder, "/analysis/estimationcode/specs/", run, "_", type, "/output", "/covar_names.csv"))
    covar_lens_csv <- read_csv(paste0(project_folder, "/analysis/estimationcode/specs/", run, "_", type, "/output", "/covar_lens.csv"))
    ests <- read_csv(paste0(project_folder, "/analysis/estimationcode/specs/", run, "_", type, "/output", "/se.csv"))
    varcov_csv <- read_csv(paste0(project_folder, "/analysis/estimationcode/specs/", run, "_", type, "/output", "/varcov.csv"))

  theta <- ests$solution
  varcov <- as.matrix(varcov_csv)
  covar_names <- covar_names_csv[['x']]
  covar_lens <- covar_lens_csv[['x']]
  cost_covar_lens <- covar_lens[1:2]
  
  # Back out individual covar names and beta params from covar_lanes
  W1_names <- covar_names[1:covar_lens[1]]
  W2_names <- covar_names[(1+covar_lens[1]):(covar_lens[1]+covar_lens[2])]
  W3_names <- covar_names[(1+covar_lens[1]+covar_lens[2]):(covar_lens[1]+covar_lens[2]+covar_lens[3])]
  Z_names  <- covar_names[(1+covar_lens[1]+covar_lens[2]+covar_lens[3]):(covar_lens[1]+covar_lens[2]+covar_lens[3]+covar_lens[4])]
  
  beta1 <- theta[1:covar_lens[1]]
  beta2 <- theta[(1+covar_lens[1]):(covar_lens[1]+covar_lens[2])]
  beta3 <- theta[(1+covar_lens[1]+covar_lens[2]):(covar_lens[1]+covar_lens[2]+covar_lens[3])]
  beta0  <- theta[(1+covar_lens[1]+covar_lens[2]+covar_lens[3]):(covar_lens[1]+covar_lens[2]+covar_lens[3]+covar_lens[4])]
  
  # Derived elambda_i
  print("Fixing median risk at 0.62, hardcoded")
  median <- 0.62
  
  varcov_alpha <- varcov[1:covar_lens[1], 1:covar_lens[1]]
  calculate_elambda <- function(W1) {
    return(1/exp(W1 %*% beta1))
  }
  calculate_elambda_se <- function(W1, elambda) {
    return(elambda * sqrt(t(W1) %*% varcov_alpha %*% W1))
  }
  # No-kid under-50
  W1 <- c(0, median, 0, 0)
  elambda_i_1 <- calculate_elambda(W1)
  se_elambda_i_1 <- calculate_elambda_se(W1, elambda_i_1)
  # Kid under-50
  W1 <- c(0, median, 1, 0)
  elambda_i_2 <- calculate_elambda(W1)
  se_elambda_i_2 <- calculate_elambda_se(W1, elambda_i_2)
  # No-kid over-50
  W1 <- c(0, median, 0, 1)
  elambda_i_3 <- calculate_elambda(W1)
  se_elambda_i_3 <- calculate_elambda_se(W1, elambda_i_3)
  # Kid over-50
  W1 <- c(0, median, 1, 1)
  elambda_i_4 <- calculate_elambda(W1)
  se_elambda_i_4 <- calculate_elambda_se(W1, elambda_i_4)
  
  # Derived omega_i
  varcov_omega <- varcov[
    (1+covar_lens[1]):(covar_lens[1]+covar_lens[2]),
    (1+covar_lens[1]):(covar_lens[1]+covar_lens[2])
  ]
  W2 <- c(1)
  omega_i <- exp(W2 %*% beta2)
  se_omega_i <- omega_i * sqrt(t(W2) %*% varcov_omega %*% W2)
  
  # Derived psi_i
  varcov_psi <- varcov[
    (1+covar_lens[1]+covar_lens[2]):(covar_lens[1]+covar_lens[2]+covar_lens[3]),
    (1+covar_lens[1]+covar_lens[2]):(covar_lens[1]+covar_lens[2]+covar_lens[3])
  ]
  W3 <- c(1)
  psi_i <- exp(W3 %*% beta3)
  se_psi_i <- psi_i * sqrt(t(W3) %*% varcov_psi %*% W3)
  
  risk_aver_i <- psi_i / 1000
  gambleloss <- log(2-exp(-risk_aver_i * 100)) / risk_aver_i
  gambleloss_delta_1 <- .1 * exp(- .1 * psi_i) / ((2-exp(- .1 * psi_i)) * psi_i)
  gambleloss_delta_2 <- log(2-exp(- .1 * psi_i)) / psi_i^2
  se_gambleloss <- 1000 * sqrt((gambleloss_delta_1-gambleloss_delta_2)^2) * se_psi_i
  
  model_results <- data.frame(
    covar_names, 
    theta[1:(sum(covar_lens[1:4]))], 
    sqrt(diag(varcov))
  )
  
  output <- data.frame(
    c("elambda_i_1", "elambda_i_2", "elambda_i_3", "elambda_i_4",
      "omega_i", "psi_i", "gambleloss"),
    c(elambda_i_1, elambda_i_2, elambda_i_3, elambda_i_4,
      omega_i, psi_i, gambleloss),
    c(se_elambda_i_1, se_elambda_i_2, se_elambda_i_3, se_elambda_i_4,
      se_omega_i, se_psi_i, se_gambleloss)
  )
  names(output) <- c("param", "est", "se")
  
  write.csv(output, paste0(output_path, "/derived.csv"))
  write.csv(model_results, paste0(output_path, "/model_results.csv"))
  
}

main()

